import os
from typing import List, Tuple

from llama_index.core.schema import TextNode

from iflytech_assistant.es import index, search
from iflytech_assistant.spark_client import instruct_with_profile as instruct

__all__ = ["generate", "generate_response", "refine"]

templates_dir = os.path.join(os.path.dirname(__file__), "templates")
template_file = os.path.join(templates_dir, "girl_code.md")
prompt_template = open(template_file).read()

INDEX_NAME = "girl_code-2024-11-11"


def build_prompt(
    user_input: str,
    examples: List[Tuple[str, str]],
) -> str:
    example_prompt = "## 举例说明\n"
    for input_, output in examples:
        example_prompt += f"{input_}\n潜台词：{output}\n\n"

    example_prompt = example_prompt.strip()

    return prompt_template.format(
        user_input=user_input,
        examples=example_prompt,
    )


def generate(user_input):
    nodes: List[TextNode] = search(user_input, index_name=INDEX_NAME, rerank=False)
    examples = [(node.text, node.metadata["answer"]) for node in nodes]
    user_prompt = build_prompt(user_input, examples)
    llm_response = instruct(user_prompt)
    return user_prompt, llm_response.llm_response


def index_exmples(examples: List[Tuple[str, str]]):
    nodes: List[TextNode] = []

    for question, answer, gender in examples:
        node = TextNode(
            text=question,
            metadata={"answer": answer, "gender": gender},
        )

        nodes.append(node)

    index(nodes, index_name=INDEX_NAME)
